""" 
    拆分文本向量为多个表格
"""

import os
import argparse
import pandas as pd


def split_table_into_parts(input, output, split_num):
    # 读取表格文件
    df = pd.read_csv(input)
    # 计算每份的行数
    total_rows = len(df)
    rows_per_part = total_rows // split_num
    remainder = total_rows % split_num
    # 确保输出目录存在
    if not os.path.exists(output):
        os.makedirs(output)
    # 拆分表格并保存
    for i in range(split_num):
        start_row = i * rows_per_part
        end_row = start_row + rows_per_part
        # 处理余数部分
        if i == split_num-1:
            end_row += remainder
        # 获取当前部分的 DataFrame
        part_df = df.iloc[start_row:end_row]
        # 保存为新的 CSV 文件
        part_file_path = os.path.join(output, f'part_{i+1}.csv')
        part_df.to_csv(part_file_path, index=False)
        
    print(f"表格已成功拆分为 { split_num } 份，保存在目录: {output}")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-input", "--input")
    parser.add_argument("-output", "--output")
    parser.add_argument("-split_num","--split_num")
    params =  parser.parse_args()
    
    return params


if __name__ == "__main__":
    params = parse_args()
    input = params.input
    output = params.output
    split_num = int(params.split_num)
    
    split_table_into_parts(input, output, split_num)